// zhiyuan-upload-books-to-dify.js
// 知源中医图书批量上传工具（ESM 兼容版）

import fs from 'fs';
import path from 'path';
import axios from 'axios';
import https from 'https';
import FormData from 'form-data';

// ========================
// 配置区
// ========================
const DIFY_KB_ID = 'd7ade38e-788a-4f7e-af6f-4891055a2597';
const AUTH_TOKEN = 'eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJ1c2VyX2lkIjoiNGZiOTJiMWItNDRiOS00OWQ5LWFlYjktYmZhZmM4MmFlYTM4IiwiZXhwIjoxNzYyMjQ1ODg2LCJpc3MiOiJTRUxGX0hPU1RFRCIsInN1YiI6IkNvbnNvbGUgQVBJIFBhc3Nwb3J0In0.LyOuEPFwMM_sYbnUPsbmOScZ6ZIYyf2BQBMD889I9kY';

const BOOKS_ROOT = path.resolve(new URL(import.meta.url).pathname, '..', 'docs', '知源中医').replace(/^\/(\w:)/, '$1'); // Windows兼容

const ALLOWED_EXT = ['.md', '.txt', '.pdf', '.docx', '.xlsx'];
const BATCH_SIZE_PER_REQUEST = 5;

// 命令行参数
const isDeleteOnly = process.argv.includes('--delete');
const isParse = process.argv.includes('--parse');
console.log(`🚀 开始上传${isParse ? '并解析' : ''}...`);
if (isDeleteOnly) console.log('⚠️  --delete 模式：仅删除所有文档');

// Axios 实例
const api = axios.create({
  baseURL: 'https://rf-dify.docker.internal:5174/api',
  headers: { Authorization: `Bearer ${AUTH_TOKEN}` },
  httpsAgent: new https.Agent({ rejectUnauthorized: false }),
});

// ========================
// 核心 API 函数
// ========================

/**
 * 获取知识库所有文档（自动翻页）
 */
async function getAllDocuments() {
  const docs = [];
  let page = 1;
  const limit = 100;
  while (true) {
    const res = await api.get(`/console/api/datasets/${DIFY_KB_ID}/documents`, {
      params: { page, limit, keyword: '' },
    });
    const data = res.data.data || [];
    docs.push(...data);
    if (!res.data.has_more) break;
    page++;
  }
  return docs;
}

/**
 * 删除指定文档
 */
async function deleteDocument(docId) {
  try {
    await api.delete(`/console/api/datasets/${DIFY_KB_ID}/documents`, {
      params: { document_id: docId },
    });
    console.log(`🗑️ 已删除: id=${docId}`);
    return true;
  } catch (err) {
    console.error(`❌ 删除失败 [${docId}]:`, err.response?.data || err.message);
    return false;
  }
}

/**
 * 上传单个文件
 */
async function uploadFile(filePath) {
  const formData = new FormData();
  formData.append('file', fs.createReadStream(filePath));
  formData.append('source', 'datasets');

  try {
    const res = await api.post('/console/api/files/upload', formData, {
      headers: {
        ...formData.getHeaders(),
        Authorization: `Bearer ${AUTH_TOKEN}`,
      },
    });
    console.log(`✅ 上传成功: ${res.data.name} → file_id=${res.data.id}`);
    return res.data.id;
  } catch (err) {
    console.error(`❌ 上传失败 [${path.basename(filePath)}]:`, err.response?.data || err.message);
    return null;
  }
}

/**
 * 创建文档（触发解析）
 */
async function createDocuments(fileIds) {
  if (fileIds.length === 0) return false;
  if (fileIds.length > 5) {
    console.error('❌ 每次最多提交5个文件');
    return false;
  }

  const payload = {
    data_source: {
      type: 'upload_file',
      info_list: {
        data_source_type: 'upload_file',
        file_info_list: {
          file_ids: fileIds,
        },
      },
    },
    indexing_technique: 'high_quality',
    process_rule: {
      mode: 'custom',
      rules: {
        pre_processing_rules: [
          { id: 'remove_extra_spaces', enabled: true },
          { id: 'remove_urls_emails', enabled: false },
        ],
        segmentation: {
          separator: '\n',
          max_tokens: 2048,
          chunk_overlap: 50,
        },
      },
    },
    doc_form: 'text_model',
    doc_language: 'Chinese Simplified',
    retrieval_model: {
      search_method: 'semantic_search',
      reranking_enable: false,
      reranking_mode: null,
      reranking_model: {
        reranking_provider_name: '',
        reranking_model_name: '',
      },
      weights: null,
      top_k: 2,
      score_threshold_enabled: false,
      score_threshold: null,
    },
    embedding_model: 'bge-m3:Q4_K_M',
    embedding_model_provider: 'langgenius/ollama/ollama',
  };

  try {
    const res = await api.post(`/console/api/datasets/${DIFY_KB_ID}/documents`, payload);
    console.log(`🚀 解析任务已提交！batch=${res.data.batch}, 文件数=${fileIds.length}`);
    return true;
  } catch (err) {
    console.error('❌ 创建文档失败:', err.response?.data || err.message);
    return false;
  }
}

/**
 * 【API】查询索引失败的文档
 */
async function getErrorDocuments() {
  try {
    const res = await api.get(`/console/api/datasets/${DIFY_KB_ID}/error-docs`);
    const errors = res.data.data || [];
    console.log(`🔍 发现 ${errors.length} 个索引失败的文档:`);
    errors.forEach(err => {
      console.log(`❌ [doc_id=${err.id}] ${err.error}`);
    });
    return errors;
  } catch (err) {
    console.error('❌ 查询错误文档失败:', err.response?.data || err.message);
    return [];
  }
}


// ========================
// 主流程
// ========================

async function main() {
  // 模式一：仅删除
  if (isDeleteOnly) {
    console.log('\n🧹 正在清空知识库...');
    const docs = await getAllDocuments();
    if (docs.length === 0) {
      console.log('🟡 知识库为空');
      return;
    }
    for (const doc of docs) {
      await deleteDocument(doc.id);
    }
    console.log('🎉 删除完成！');
    return;
  }

  // 获取现有文档（用于查重）
  const existingDocs = await getAllDocuments();
  const existingNames = new Set(existingDocs.map(d => d.name));

  // 获取书籍目录列表
  if (!fs.existsSync(BOOKS_ROOT)) {
    console.error(`❌ 图书根目录不存在: ${BOOKS_ROOT}`);
    process.exit(1);
  }

  const bookDirs = fs.readdirSync(BOOKS_ROOT, { withFileTypes: true })
    .filter(d => d.isDirectory())
    .map(d => d.name);

  if (bookDirs.length === 0) {
    console.log('🟡 未发现任何书籍目录');
    return;
  }

  // 遍历每本书，上传后立即尝试解析
  for (const bookName of bookDirs) {
    const bookPath = path.join(BOOKS_ROOT, bookName);

    // 获取该书的所有文件
    const files = fs.readdirSync(bookPath)
      .filter(f => ALLOWED_EXT.includes(path.extname(f).toLowerCase()))
      .map(f => path.join(bookPath, f));

    if (files.length === 0) {
      console.log(`🟡 [${bookName}] 无支持格式文件，跳过`);
      continue;
    }

    console.log(`\n📘 正在处理书籍: ${bookName} (${files.length} 个文件)`);

    const fileIdsForThisBook = [];

    for (const filePath of files) {
      const filename = path.basename(filePath);
      const docName = `${bookName} - ${filename}`;

      if (existingNames.has(docName)) {
        console.log(`🔁 跳过已存在: ${docName}`);
        continue;
      }

      const fileId = await uploadFile(filePath);
      if (fileId) {
        fileIdsForThisBook.push(fileId);
      }
    }

    // 🚀 当前这本书上传完毕，立即提交解析（如果启用了 --parse）
    if (isParse && fileIdsForThisBook.length > 0) {
      console.log(`🧠 "${bookName}" 上传完成，正在提交解析任务...`);

      // 分批提交（每5个一组）
      for (let i = 0; i < fileIdsForThisBook.length; i += BATCH_SIZE_PER_REQUEST) {
        const batch = fileIdsForThisBook.slice(i, i + BATCH_SIZE_PER_REQUEST);
        await createDocuments(batch);
      }
    } else if (fileIdsForThisBook.length > 0) {
      console.log(`📌 "${bookName}" 上传完成！共 ${fileIdsForThisBook.length} 个新文件。使用 --parse 可触发解析。`);
    } else {
      console.log(`🟢 "${bookName}" 所有文件均已存在，无需上传`);
    }
    break
  }

  // 检查错误
  await getErrorDocuments();

  console.log('\n🎉 全部任务完成！');
}

// 启动主函数
main().catch(console.error);